# /home_ssd/lhc/hand_detect_v3/2025-10-12_11_47_02 : 0.0 => 1.0
# /home_ssd/lhc/hand_detect_v3/2025-10-12_11_51_25 : 0.0 => 2.0
# /home_ssd/lhc/hand_detect_v3/2025-10-12_11_52_24 : 0.0 => 2.0
# /home_ssd/lhc/hand_detect_v3/2025-10-12_11_53_08 : 0.0 => 3.0
# /home_ssd/lhc/hand_detect_v3/2025-10-12_11_53_46 : 0.0 => 3.0
# /home_ssd/lhc/hand_detect_v3/2025-10-12_11_54_38 : 0.0 => 4.0
# /home_ssd/lhc/hand_detect_v3/2025-10-12_11_55_13 : 0.0 => 4.0
# /home_ssd/lhc/hand_detect_v3/2025-10-12_11_55_56 : 0.0 => 1.0
# /home_ssd/lhc/hand_detect_v3/2025-10-12_11_56_29 : 0.0 => 1.0
# /home_ssd/lhc/hand_detect_v3/2025-10-12_12_02_38 : 0.0 => 1.0



import os

path = r"/home_ssd/lhc/hand_detect_v3"

eve_txt = "event_frame_window_new.txt"
img_txt = "rgb_wrap_window_new.txt"



# 内容
'''
0_0010.txt: 4 0.373828 0.538194 0.153906 0.170833
0_0011.txt: 4 0.394141 0.519444 0.153906 0.180556
0_0012.txt: 
0_0013.txt: 4 0.434375 0.490972 0.160938 0.198611
0_0014.txt: 
0_0015.txt: 4 0.477734 0.475000 0.158594 0.205556
0_0016.txt: 4 0.499219 0.462500 0.162500 0.205556
0_0017.txt: 4 0.515625 0.440972 0.159375 0.201389
'''


import re
import shutil
import argparse


def parse_corrections_from_comments(script_path):
	"""从脚本开头的注释中解析修正线索，格式示例：
	# /home.../2025-10-12_11_47_02 : 0.0 => 1.0
	返回字典：{ '2025-10-12_11_47_02': (from_class, to_class), ... }
	"""
	pattern = re.compile(r'#\s*(?P<fullpath>.+?):\s*(?P<from>[0-9]+(?:\.[0-9]+)?)\s*=>\s*(?P<to>[0-9]+)')
	corrections = {}
	try:
		with open(script_path, 'r', encoding='utf-8') as f:
			for line in f:
				m = pattern.search(line)
				if m:
					full = m.group('fullpath').strip()
					from_c = float(m.group('from'))
					to_c = int(m.group('to'))
					key = os.path.basename(full)
					corrections[key] = (from_c, to_c)
	except Exception:
		pass
	return corrections


def process_annotation_line(line, from_c, to_c):
	"""处理单行注释格式：
	例子： '0_0010.txt: 4 0.373828 0.538194 0.153906 0.170833' 或 '0_0012.txt:'
	如果行包含类别且类别等于 from_c，则替换为 to_c（整数），返回新的行与是否修改标志。
	"""
	# 保留原行结尾
	newline = line.rstrip('\n')
	if ':' not in newline:
		return line, False
	prefix, rest = newline.split(':', 1)
	rest_strip = rest.strip()
	if rest_strip == '':
		# 无目标，保持不变
		return line, False
	# 取第一个 token 为 class
	tokens = rest_strip.split()
	try:
		cls_val = float(tokens[0])
	except Exception:
		return line, False
	if abs(cls_val - from_c) < 1e-6:
		# 替换为整数 to_c
		tokens[0] = str(int(to_c))
		new_rest = ' '.join(tokens)
		newline2 = f"{prefix}: {new_rest}\n"
		return newline2, True
	else:
		return line, False


def process_txt_file(file_path, from_c, to_c, dry_run=True, backup=True):
	"""处理单个 txt 文件，逐行替换类别。返回(修改次数, 文件路径)。"""
	if not os.path.exists(file_path):
		return 0, file_path
	with open(file_path, 'r', encoding='utf-8') as f:
		lines = f.readlines()

	changed = 0
	new_lines = []
	for line in lines:
		nl, modified = process_annotation_line(line, from_c, to_c)
		if modified:
			changed += 1
		new_lines.append(nl)

	if changed and not dry_run:
		if backup:
			shutil.copyfile(file_path, file_path + '.bak')
		with open(file_path, 'w', encoding='utf-8') as f:
			f.writelines(new_lines)

	return changed, file_path


def fix_subfolders(root_path, eve_txt_name=eve_txt, img_txt_name=img_txt, corrections=None, dry_run=True):
	"""遍历 root_path 下的每个子文件夹，根据 corrections（dict: basename -> (from,to)）修正两个 txt 文件中的类别。
	返回修改汇总字典。
	"""
	if corrections is None:
		corrections = {}
	summary = {}
	for name in os.listdir(root_path):
		sub = os.path.join(root_path, name)
		if not os.path.isdir(sub):
			continue
		key = name
		if key not in corrections:
			continue
		from_c, to_c = corrections[key]
		eve_path = os.path.join(sub, eve_txt_name)
		img_path = os.path.join(sub, img_txt_name)
		eve_changed, _ = process_txt_file(eve_path, from_c, to_c, dry_run=dry_run)
		img_changed, _ = process_txt_file(img_path, from_c, to_c, dry_run=dry_run)
		if eve_changed or img_changed:
			summary[key] = {'from': from_c, 'to': to_c, 'eve_changed': eve_changed, 'img_changed': img_changed, 'path': sub}
	return summary


def main():
	parser = argparse.ArgumentParser(description='修正子文件夹中两个 txt 的类别标注（根据脚本顶部注释或手动传入映射）')
	parser.add_argument('--path', type=str, default=path, help='根目录，包含多个子文件夹')
	parser.add_argument('--apply', action='store_true', help='实际写回文件（否则为 dry-run）')
	parser.add_argument('--script', type=str, default=__file__, help='用于解析注释中修正线索的脚本路径')
	args = parser.parse_args()

	corrections = parse_corrections_from_comments(args.script)
	if not corrections:
		print('未在脚本注释中发现修正线索。请在脚本顶部按照示例添加，或手动调用函数。')
	else:
		print('解析到修正线索：')
		for k, v in corrections.items():
			print(f"  {k}: {v[0]} => {v[1]}")

	dry_run = not args.apply
	summary = fix_subfolders(args.path, corrections=corrections, dry_run=dry_run)
	if not summary:
		print('未发现需修正的子文件夹或修正后无变更。')
	else:
		print('修正摘要：')
		for k, v in summary.items():
			print(f"  {k}: from {v['from']} to {v['to']}, eve_changed={v['eve_changed']}, img_changed={v['img_changed']}, path={v['path']}")
		if dry_run:
			print('\n注意：当前为 dry-run（未写回文件）。若确认无误，请使用 --apply 参数执行写回。')


if __name__ == '__main__':
	main()



